package com.piece.module.generator.support;

import com.piece.core.framework.util.string.StringUtil;
import com.sun.codemodel.CodeWriter;
import com.sun.codemodel.JPackage;
import org.apache.commons.lang.StringEscapeUtils;

import java.io.ByteArrayOutputStream;
import java.io.FilterOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

public class PrintCodeWriter extends CodeWriter {

    private final Map<String, ByteArrayOutputStream> streams;
    private final Map<String, String> contents;

    public PrintCodeWriter() {
        this.streams = new HashMap<>();
        this.contents = new HashMap<>();
    }

    @Override
    public OutputStream openBinary(JPackage jPackage, String s) throws IOException {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        this.streams.put(jPackage.name() + "." + s, outputStream);
        return new FilterOutputStream(outputStream) {
            public void close() {
                // don't let this stream close
            }
        };
    }

    public Map<String, String> contents() {
        return this.contents;
    }

    @Override
    public void close() throws IOException {
        streams.entrySet().stream().forEach(entry -> {
            ByteArrayOutputStream stream = entry.getValue();
            String s = StringEscapeUtils.unescapeJava(StringUtil.toEncodedString(stream.toByteArray(), StandardCharsets.UTF_8));
            String text = StringUtil.trim(s);
            contents.put(entry.getKey(), text);
            try {
                stream.close();
            } catch (Exception e) {
                stream = null;
            }
        });
    }
}
